from typing import List, Optional, Dict
from pydantic import BaseModel, Field
from config import Config

# 数据模型
class QueryRequest(BaseModel):
    query_text: str = Field(..., description="查询文本")
    query_type: str = Field(..., description="查询类型：功能用户需求、功能过程、子过程描述")
    cosmicid: Optional[str] = Field(None, description="cosmic ID过滤")
    contract_code: Optional[str] = Field(None, description="合同编码过滤")
    firmcode: Optional[str] = Field(None, description="厂商编码过滤")
    exclude_pswfinstid: Optional[str] = Field(None, description="排除的pswfinstid")
    top_k: int = Field(Config.DEFAULT_TOP_K, ge=1, le=Config.MAX_TOP_K, description="返回top K个相似结果")


class BatchQueryRequest(BaseModel):
    queries: List[QueryRequest] = Field(..., description="批量查询请求")


class UpdateRequest(BaseModel):
    data_id: str = Field(..., description="数据ID")
    cosmicid: str = Field(..., description="cosmic ID")
    firmcode: str = Field(..., description="厂商编码")
    contract_code: str = Field(..., description="合同编码")
    gnyhxq: Optional[str] = Field(None, description="功能用户需求")
    gngc: Optional[str] = Field(None, description="功能过程")
    zgcms: Optional[str] = Field(None, description="子过程描述")


class BatchUpdateRequest(BaseModel):
    updates: List[UpdateRequest] = Field(..., description="批量更新请求")


class DeleteRequest(BaseModel):
    data_ids: List[str] = Field(..., description="要删除的cosmicid列表")


class SimilarityResult(BaseModel):
    data_id: str
    cosmicid: str
    firmcode: str
    contract_code: str
    content: str
    similarity_score: float
    similarity_percentage: float


class QueryResponse(BaseModel):
    query_text: str
    query_type: str
    results: List[SimilarityResult]
    total_count: int
    execution_time: float


class ThresholdUpdateRequest(BaseModel):
    thresholds: Dict[str, float] = Field(..., description="相似度阈值配置")
    save_to_file: bool = Field(True, description="是否保存到配置文件")

    class Config:
        schema_extra = {
            "example": {
                "thresholds": {
                    "功能用户需求": 0.85,
                    "功能过程": 0.90,
                    "子过程描述": 0.88
                },
                "save_to_file": True
            }
        }


class MultiFieldSimilarityRequest(BaseModel):
    cosmicid: str = Field(..., description="cosmic ID")
    fields: Optional[List[str]] = Field(None, description="要分析的字段列表，默认为['gnyhxq', 'gngc', 'zgcms']")

    class Config:
        schema_extra = {
            "example": {
                "cosmicid": "COSMIC123456",
                "fields": ["gnyhxq", "gngc", "zgcms"]
            }
        }


class PswfinstidSimilarityRequest(BaseModel):
    pswfinstid: str = Field(..., description="工作流实例ID")

    class Config:
        schema_extra = {
            "example": {
                "pswfinstid": "WF123456789"
            }
        }